# src/tcrun.py

import os
import re
import sys
import shutil
import logging
import tomllib
import traceback
import concurrent.futures
from typing import Generator

import pandas as pd

from src.vmctrl import VmwareController
from src.macro import AlloySystem, generate_macro
from src.ctrl import is_stopped

logger = logging.getLogger(__name__)


class TcRunner:
    def __init__(self):
        self.config = {}
        self.vmctrl: VmwareController = None
        self.required_keys = ["vmware", "thermo_calc", "alloy_system", "output"]
        self.vm_req_keys = ["vmx", "vmrun", "vm_user", "vm_passwd"]
        self.tc_req_keys = ["tcp_path", "work_root"]
        self.alloy_req_keys = ["comps_file", "tdb_file"]
        self.out_req_keys = ["output_root"]

    def load_config(self, config_path: str) -> dict:
        """Load configuration file"""
        if not os.path.exists(config_path):
            err = f"No Found config({config_path})"
            logger.error(err)
            raise FileNotFoundError(err)
        with open(config_path, "rb") as f:
            return tomllib.load(f)

    def check_config(self, data: dict, req_keys: list[str]):
        """Check if data missing required keys."""
        missing_keys = [k for k in req_keys if k not in data]
        if missing_keys:
            raise KeyError(
                f"Config missing items: {missing_keys} / {list(data.keys())}"
            )

    def load_compositions(self, comps_file: str):
        if not os.path.exists(comps_file):
            raise FileNotFoundError(f"No found composition_data({comps_file})")

        # When reading file, use the first row as column names
        if comps_file.endswith(".csv") or comps_file.endswith(".CSV"):
            df = pd.read_csv(comps_file)
        elif comps_file.endswith(".xlsx") or comps_file.endswith(".xls"):
            df = pd.read_excel(comps_file)
        else:
            raise ValueError(f"Unsupported composition data file format: {comps_file}")

        # Validate column name format
        if len(df.columns) < 2 or df.columns[0].lower() != "phases":
            raise ValueError(f"Data file format error: First column must be 'phases'")

        pattern = r"([A-Za-z]+)\_?((?:\d+(?:\.\d*)?|\.\d+)(?:[Ee][+-]?\d+)?)?"
        for _, row in df.iterrows():
            # Skip empty rows
            if pd.isna(row["phases"]):
                continue

            # Get phase name list
            phases = sorted(p.lower() for p in str(row["phases"]).split())

            # Process composition data
            names = []
            data = {}
            for cell in row[1:]:  # Skip phases column
                if pd.isna(cell):
                    continue

                cell_str = str(cell).strip()
                match = re.match(pattern, cell_str)
                if not match:
                    continue
                names.append(cell_str)
                comp, fract = match.groups()
                data[comp] = 1.0 if fract is None else float(fract)

            if not data:
                continue

            # Normalize, round to 6 decimal places
            total = sum(data.values())
            normalized = {e: round(f / total, 6) for e, f in data.items()}
            yield phases, names, normalized

    def make_workdir(
        self,
        vm_dir: str,
        host_tdb_file: str,
        phases: list[str],
        names: list[str],
        alloy_system: AlloySystem,
        vmctrl: VmwareController = None,
    ) -> tuple[str, str]:
        """Create working directory and copy TDB and TCM.

        Args:
            vm_dir: Virtual machine directory
            host_tdb_file: Host TDB file
            phases: Phase list
            names: Component name list
            alloy_system: Alloy system
            vmctrl: Virtual machine controller
        Returns:
            Tuple:
                tuple(work_dir, vm_tcm_file)
        """
        if vmctrl is None:
            vmctrl = self.vmctrl

        name = "-".join(sorted(phases) + sorted(names))
        work_dir = os.path.join(vm_dir, name)

        # Create working directory
        if vmctrl.is_dir_exists(work_dir):
            logger.info(f"Virtual machine directory {work_dir} already exists")
            vmctrl.delete_dir(work_dir)
        vmctrl.create_dir(work_dir)

        # Copy TDB file
        tdb_filename = os.path.basename(host_tdb_file)
        vm_tdb_file = os.path.join(work_dir, tdb_filename)
        result = vmctrl.copy_file_from_host_to_guest(host_tdb_file, vm_tdb_file)
        logger.debug(result)
        logger.info(f"Copied TDB file to virtual machine: {vm_tdb_file}")

        # Generate TCM file
        tcm_filename = f"tcp_macro.TCM"
        vm_tcm_file = os.path.join(work_dir, tcm_filename)
        tcm_str = generate_macro(tdb_filename, alloy_system)
        result = vmctrl.write_file(vm_tcm_file, tcm_str, crlf=True)
        logger.debug(result)
        logger.info(f"Generated TCM file and copied to virtual machine: {vm_tcm_file}")

        return work_dir, vm_tcm_file

    def run_tcp_macro(
        self, tcp_path: str, vm_tcm_file: str, vmctrl: VmwareController = None
    ):
        """Run TCP program"""
        if vmctrl is None:
            vmctrl = self.vmctrl

        logger.info(f"Run program {tcp_path} {vm_tcm_file}")
        result = vmctrl.run_program(tcp_path, vm_tcm_file)
        logger.debug(result)

        vm_poly3_path = os.path.join(os.path.dirname(vm_tcm_file), "RESULT.POLY3")
        return vmctrl.is_file_exists(vm_poly3_path)

    def copy_result(
        self, work_dir: str, host_dir: str, vmctrl: VmwareController = None
    ):
        """Copy result files to host"""
        if vmctrl is None:
            vmctrl = self.vmctrl

        out_dir = os.path.join(host_dir, os.path.basename(work_dir))
        os.makedirs(host_dir, exist_ok=True)
        if os.path.exists(out_dir):
            shutil.rmtree(out_dir)
        os.makedirs(out_dir)
        result = vmctrl.copy_dir_from_guest_to_host(work_dir, out_dir)
        logger.debug(result)
        logger.info(f"Copied result files to host: {out_dir}")
        return out_dir

    def cleanup(self, work_dir: str, vmctrl: VmwareController = None):
        """Clean up working directory"""
        if work_dir and vmctrl:
            vmctrl.delete_dir(work_dir)
            logger.info(f"Cleaned up working directory: {work_dir}")

    def should_stop(self) -> bool:
        """Stop processing"""
        if is_stopped():
            logger.info("Interrupt signal received, canceling processing...")
            return True
        return False

    def process_composition(
        self,
        vm_dir: str,
        tcp_path: str,
        host_dir: str,
        tdb_file: str,
        phases: list[str],
        names: list[str],
        data: dict,
    ) -> str | None:
        """Process a single composition combination"""
        vmctrl = VmwareController(**self.config["vmware"])
        alloy_system = AlloySystem(
            phases=phases, comps=list(data.keys()), fracts=list(data.values())
        )

        try:
            if self.should_stop():
                return None
            work_dir, vm_tcm_file = self.make_workdir(
                vm_dir, tdb_file, phases, names, alloy_system, vmctrl
            )

            if self.should_stop():
                return None
            res = self.run_tcp_macro(tcp_path, vm_tcm_file, vmctrl)

            if self.should_stop():
                return None
            elif not res:
                logger.error(f"Failed to run program: {work_dir}")
                return None
            else:
                out_dir = self.copy_result(work_dir, host_dir, vmctrl)
                return out_dir

        except Exception as e:
            logger.error(f"Error processing composition combination: {e}")
            logger.error(traceback.format_exc())
        finally:
            self.cleanup(work_dir, vmctrl)

    def run(self, config_path: str, parallel: int = 1) -> Generator[str, None, None]:
        """Run (synchronous)

        parallel: Number of parallel workers
        """
        # Load and check configuration
        self.config = self.load_config(config_path)
        self.check_config(self.config, self.required_keys)
        self.check_config(self.config["vmware"], self.vm_req_keys)
        self.check_config(self.config["thermo_calc"], self.tc_req_keys)
        self.check_config(self.config["alloy_system"], self.alloy_req_keys)
        self.check_config(self.config["output"], self.out_req_keys)

        self.vmctrl = VmwareController(**self.config["vmware"])
        tcp_path = self.config["thermo_calc"]["tcp_path"]
        vm_dir = self.config["thermo_calc"]["work_root"]
        tdb_file = self.config["alloy_system"]["tdb_file"]
        comps_file = self.config["alloy_system"]["comps_file"]
        host_dir = self.config["output"]["output_root"]

        results = []

        logger.info(f"Use {parallel} worker processes for parallel computation")

        # Use thread pool to execute parallel tasks
        with concurrent.futures.ThreadPoolExecutor(max_workers=parallel) as executor:
            futures = []

            for phases, names, data in self.load_compositions(comps_file):
                if self.should_stop():
                    break

                # Submit task to thread pool
                future = executor.submit(
                    self.process_composition,
                    vm_dir,
                    tcp_path,
                    host_dir,
                    tdb_file,
                    phases,
                    names,
                    data,
                )
                futures.append((future, phases, names))

            # Collect results
            for future, phases, names in futures:
                if self.should_stop():
                    break
                try:
                    res = future.result()
                    if res is not None:
                        yield res
                    results.append(res)
                except Exception as e:
                    logger.error(
                        f"Exception occurred processing composition {phases} {names}: {e}"
                    )
                    results.append(None)

        successes = [res for res in results if res is not None]
        logger.info(f"Successfully processed {len(successes)} / {len(results)}")
